import gym_super_mario_bros
from nes_py.wrappers import JoypadSpace
from gym_super_mario_bros.actions import SIMPLE_MOVEMENT


def gymnasium_reset(self, **kwargs):
    return self.env.reset(), {}

def make_super_mario_env():
    env = gym_super_mario_bros.make('SuperMarioBros-v0')
    env = JoypadSpace(env, SIMPLE_MOVEMENT)  # set the joypad wrapper
    return env

env = make_super_mario_env()

done = True
for step in range(50000):
    if done:
        state, info = env.reset()
    state, reward, done, trunc, info = env.step(env.action_space.sample())
    env.render()


env.close()